import logging
import os
import time

import torch

from _dynamic_dataset.AccCombinationRecord import WriteRecordsToFile
from server_client import copy_model_params
from server_client.Server import Server
from utils.datasetUtils import get_DDSS_clients
from utils.modelUtils import getResNet18


def Algorithm_True_Shapely(num_epochs, clients, server, device, combinations, dirPath, acc_record):
    os.makedirs(dirPath, exist_ok=True)
    logging.basicConfig(filename=f'{dirPath}/main.log', level=logging.INFO)
    # 开始时间
    start_time = time.time()
    for combination in combinations:
        logging.info(f"当前验证的组合为{combination}")
        # 初始化服务端
        # 创建相同分布相同大小的客户端
        model = getResNet18()
        for i in range(len(clients)):
            copy_model_params(clients[i].local_model, model)
        copy_model_params(server.global_model, model)
        copy_model_params(server.sub_model, model)
        copy_model_params(server.old_model, model)
        print(f"服务器初始化成功！客户端组合为{combination}")
        # 创建客户端子集
        com_clients = []
        for clientId in combination:
            com_clients.append(clients[clientId])
        # 开始训练
        for epoch in range(1, num_epochs + 1):
            # 客户端训练
            for client in com_clients:
                client.train(server.global_model)
            # 在服务端进行集合
            # 聚合模型
            server.model_aggregate(com_clients, server.global_model)
            # 在测试集上进行测试
            correct = 0
            total = 0
            with torch.no_grad():
                for data in server.eval_loader:
                    images, labels = data
                    images, labels = images.to(device), labels.to(device)
                    outputs = server.global_model(images)
                    _, predicted = torch.max(outputs.data, 1)
                    total += labels.size(0)
                    correct += (predicted == labels).sum().item()
            # 真实的正确率
            acc = 100 * correct / total
            acc_record.addCombinationAcc(epoch, combination, acc)
    # 结束时间
    end_time = time.time()
    logging.info(f"总耗时为{end_time - start_time}s")
    WriteRecordsToFile(dirPath + "/acc_records.json", acc_record)



